package com.example.mcpclient.controller;

import com.example.mcpclient.request.InputMsgRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author xuguang
 */
@RequestMapping("/test")
@Slf4j
@RestController
public class TestController {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ApplicationContext applicationContext;


    /**
     * 使用内部的tools, 而不是mcp下游的tools
     */
    @PostMapping("/innerTools")
    public Mono<String> innerTools(@RequestBody InputMsgRequest userInput) {
        return Mono.fromCallable(() -> chatClient
                .prompt(userInput.getQuestion())
                .call()
                .content()).subscribeOn(Schedulers.boundedElastic());
    }


    /**
     * 通过rag的向量来获取, 不通过mcp的工具
     */
    @PostMapping("/rag")
    public Mono<String> rag(@RequestBody InputMsgRequest userInput) {
        return Mono.fromCallable(() -> chatClient
                .prompt(userInput.getQuestion())
                .call()
                .content()).subscribeOn(Schedulers.boundedElastic());
    }


    /**
     * 直接返回结果
     */
    @PostMapping("/input")
    public String input(@RequestBody InputMsgRequest userInput) {
        return chatClient
                .prompt(userInput.getQuestion())
                .call()
                .content(); // 流式返回内容片段
    }


    /**
     * 返回流式结果
     */
    @PostMapping("/input-stream")
    public Flux<String> testStream(@RequestBody InputMsgRequest userInput) {
        return chatClient
                .prompt(userInput.getQuestion())
                .stream()
                .content(); // 流式返回内容片段
    }

    /**
     * 获取有哪些mcp工具
     */
    @PostMapping("/tools")
    public Mono<List<Map<String, String>>> getAvailableTools() {
        return Mono.defer(() -> {
            List<Map<String, String>> tools = new ArrayList<>();
            Map<String, ToolCallbackProvider> beansOfType = applicationContext.getBeansOfType(ToolCallbackProvider.class);
            for (ToolCallbackProvider toolProvider : beansOfType.values()) {
                ToolCallback[] callbacks = toolProvider.getToolCallbacks();
                for (ToolCallback callback : callbacks) {
                    Map<String, String> toolInfo = new HashMap<>();
                    toolInfo.put("name", callback.getToolDefinition().name());
                    toolInfo.put("description", callback.getToolDefinition().description());
                    tools.add(toolInfo);
                }
            }
            return Mono.just(tools);
        }).subscribeOn(Schedulers.boundedElastic());
    }

}
